'''
@create_time: 2025/10/23 下午4:53
@Author: GeChao
@File: db_connection.py
'''
import sqlite3
from tools.log import Logger
import pymysql
from abc import ABC, abstractmethod
from typing import Optional, List, Type, Dict, Tuple, Any, Union
from pymysql import IntegrityError, OperationalError, DatabaseError
from typing import Union, Any

logger = Logger().logger
MailHostIp = "smtp.exmail.qq.com"
MailHostPort = 465
MailSendUser = "ops@mabangerp.com"
MailSendPasswd = "Abcd.1234"


class BaseConnector(ABC):

    def __init__(self, **kwargs):
        pass

    def connect(self):
        return self

    def select(self, table: str, fields: List[str], where: Optional[Dict[str, Any]] = None, limit: Optional[int] = None) -> \
            Union[tuple[tuple[
                Any, ...], ...], int]:
        ...

    def insert(self, *args, **kwargs):
        ...

    def update(self, *args, **kwargs):
        ...

    def delete(self, table: str, where: Optional[Dict[str, Any]], commit: bool = True) -> int:
        ...

    def execute(self, *args, **kwargs):
        ...

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass


class SqlLiteConnector(BaseConnector):
    """SQLite数据库连接器，CRUD操作及事务管理
    """

    def __init__(self, db_path: str):
        """
        Args:
            db_path: 数据库文件路径
        """
        super().__init__()
        self.db_path = db_path
        self.conn: Optional[sqlite3.Connection] = None
        self.cursor: Optional[sqlite3.Cursor] = None

    def connect(self):
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        # 启用外键约束（SQLite默认关闭）
        self.cursor.execute("PRAGMA foreign_keys = ON")
        return self

    def execute(
            self,
            sql: str,
            params: Optional[Union[list, tuple, dict]] = None,
            commit: bool = True
    ) -> Union[list, int, None]:
        """执行SQL语句的统一入口
        Args:
            sql: 含占位符的SQL语句（使用?作为参数占位符）
            params: 查询参数，支持列表/元组/字典格式
            commit: 是否自动提交事务

        Returns:
            - SELECT查询：返回结果列表
            - INSERT操作：返回最后插入行ID（单条）或ID列表（批量）
            - UPDATE/DELETE：返回影响行数

        Raises:
            DatabaseError: 数据库操作异常时抛出
        """
        try:
            if not self.conn:
                self.__enter__()

            if params:
                self.cursor.execute(sql, params)
            else:
                self.cursor.execute(sql)

            if commit and self.conn:
                self.conn.commit()

                # 处理返回结果
            if sql.strip().upper().startswith("SELECT"):
                return self.cursor.fetchall()
            elif sql.strip().upper().startswith("INSERT"):
                return self.cursor.lastrowid
            else:
                return self.cursor.rowcount

        except sqlite3.Error as e:
            if self.conn:
                self.conn.rollback()
            logger.error(f"SQLite执行失败: {e}\nSQL: {sql}")
            raise DatabaseError(f"数据库操作异常: {e}")

    def select(
            self,
            table: str,
            fields: List[str],
            where: Optional[Dict[str, Any]] = None,
            limit: Optional[int] = None
    ) -> list:
        """执行查询操作

        Args:
            table: 目标表名称
            fields: 查询字段列表
            where: 查询条件字典（字段名: 期望值）
            limit: 结果集最大行数限制

        Returns:
            包含查询结果的列表，每行数据以字典形式存储
        """

        field_list = ', '.join(fields)
        sql = f"SELECT {field_list} FROM {table}"
        params = []

        if where:
            conditions = []
            for k, v in where.items():
                conditions.append(f"{k} = ?")
                params.append(v)
            sql += " WHERE " + " AND ".join(conditions)

        if limit:
            sql += f" LIMIT ?"
            params.append(limit)

        return self.execute(sql, params)

    def insert(
            self,
            table: str,
            data: Union[Dict[str, Any], List[Dict[str, Any]]],
            replace: bool = False,
            commit: bool = True
    ) -> Union[int, List[int]]:
        """执行插入或替换操作
        Args:
            table: 目标表名称
            data: 单条数据字典或批量数据字典列表
            replace: 是否启用REPLACE INTO替换已有数据
            commit: 是否自动提交事务

        Returns:
            - 单条插入：返回最后插入行ID (int)
            - 批量插入：返回插入行ID列表 (List[int])
        """
        operation = "REPLACE" if replace else "INSERT"

        if isinstance(data, list) and not data:
            return []

        items = data if isinstance(data, list) else [data]
        columns = items[0].keys()

        # 构建批量参数
        params = [tuple(item.values()) for item in items]
        placeholders = ', '.join(['?'] * len(columns))

        sql = f"{operation} INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"

        # 执行批量操作
        if len(items) > 1:
            self.cursor.executemany(sql, params)
        else:
            self.cursor.execute(sql, params[0])
        if commit and self.conn:
            self.conn.commit()

            # 获取插入ID
        if len(items) == 1:
            return self.cursor.lastrowid
        else:
            return [self.cursor.lastrowid - i for i in reversed(range(len(items)))]

    def update(
            self,
            table: str,
            updates: Dict[str, Any],
            where: Dict[str, Any],
            commit: bool = True
    ) -> int:
        """执行安全更新操作

        Args:
            table: 目标表名称
            updates: 更新字段字典（字段名: 新值）
            where: 条件字典（字段名: 期望值）
            commit: 是否自动提交事务

        Returns:
            本次更新影响的行数
        """
        set_clause = ', '.join([f"{k} = ?" for k in updates])
        where_clause = ' AND '.join([f"{k} = ?" for k in where])

        sql = f"UPDATE {table} SET {set_clause} WHERE {where_clause}"
        params = list(updates.values()) + list(where.values())

        return self.execute(sql, params, commit)

    def delete(
            self,
            table: str,
            where: Optional[Dict[str, Any]],
            commit: bool = True
    ) -> int:
        """执行安全删除操作

        Args:
            table: 目标表名称
            where: 条件字典（字段名: 期望值）
            commit: 是否自动提交事务

        Returns:
            本次删除影响的行数
        """
        if where is None:
            sql = f"DELETE FROM {table}"
            return self.execute(sql)

        else:
            where_clause = ' AND '.join([f"{k} = ?" for k in where])
            sql = f"DELETE FROM {table} WHERE {where_clause}"

            return self.execute(sql, list(where.values()), commit)

    def __enter__(self):
        self.conn = sqlite3.connect(self.db_path)
        self.cursor = self.conn.cursor()
        # 启用外键约束（SQLite默认关闭）
        self.cursor.execute("PRAGMA foreign_keys = ON")
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if self.cursor:
            self.cursor.close()
        if self.conn:
            self.conn.close()


class MysqlConnector(BaseConnector):
    """MySQL数据库连接器
    """

    def __init__(self, host=None, port=None, usr=None, pwd=None, db=None):
        """
        Args:
            host: 数据库服务器地址（默认localhost）
            port: 端口号（默认3306）
            usr: 数据库用户名
            pwd: 数据库密码
            db: 目标数据库名称
        """
        super().__init__()
        self.host = host
        self.port = port
        self.user = usr
        self.password = pwd
        self.db = db

    def connect(self):
        """
            显示链接数据库
        """
        self.conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password
                                    , db=self.db, charset='utf8')
        self.cursor = self.conn.cursor()
        return self

    def execute(
            self,
            sql: str,
            params: Optional[Union[list, tuple, dict]] = None,
            commit: bool = True
    ) -> Union[tuple[tuple[Any, ...], ...], int]:
        """执行SQL语句的统一入口

        Args:
            sql: 含占位符的SQL语句（使用%s或命名参数）
            params: 查询参数，支持列表/元组/字典格式
            commit: 是否自动提交事务

        Returns:
            - SELECT查询：返回结果元组（每行数据为嵌套元组）
            - INSERT操作：返回最后插入行ID（单条）或影响行数（批量）
            - UPDATE/DELETE：返回影响行数

        Raises:
            DatabaseError: 发生完整性错误或操作异常时抛出
        """
        try:

            self.cursor.execute(sql, params)
            if commit:
                self.conn.commit()
            if self.cursor.description:  # SELECT
                return self.cursor.fetchall()
            elif self.cursor.lastrowid:  # INSERT
                return self.cursor.lastrowid
            else:  # UPDATE/DELETE
                return self.cursor.rowcount
        except Exception as e:
            self.conn.rollback()
            logger.error(f"SQL执行失败: {e}\nSQL: {sql}")
            raise DatabaseError(f"数据库操作异常: {e}")

    def select(self, table: str, fields: List[str], where: Optional[Dict[str, Any]] = None, limit: Optional[int] = None) -> \
            Union[tuple[tuple[
                Any, ...], ...], int]:
        """执行查询操作

        Args:
            table: 目标表名称
            fields: 查询字段列表
            where: 查询条件字典（字段名: 期望值）

        Returns:
            包含查询结果的元组，每行数据以元组形式存储
        """
        field_list = ', '.join(fields)
        sql = f"SELECT {field_list} FROM {table}"
        params = []

        if where:
            conditions = []
            for k, v in where.items():
                conditions.append(f"{k} = %s")
                params.append(v)
            sql += " WHERE " + " AND ".join(conditions)

        if limit:
            sql += f" LIMIT %s"
            params.append(limit)
        return self.execute(sql, params)

    def insert(self, table: str, data: Union[Dict[str, Any], List[Dict[str, Any]]], ignore_duplicate: bool = False,
               commit: bool = True) -> Union[int, List[int]]:
        """执行插入操作（支持批量）

        Args:
            table: 目标表名称
            data: 单条数据字典或批量数据字典列表
            ignore_duplicate: 是否启用INSERT IGNORE语法
            commit: 是否自动提交事务

        Returns:
            - 单条插入：返回最后插入行ID (int)
            - 批量插入：返回首条插入行ID (int)

        Note:
            批量插入时实际返回首个插入ID，后续ID为连续自增
        """
        if not data:
            return [] if isinstance(data, list) else None

        is_bulk = isinstance(data, list)
        items = data if is_bulk else [data]

        # 构建参数化查询
        columns = ', '.join(items[0].keys())
        placeholders = ', '.join(['%s'] * len(items[0]))
        values = [tuple(item.values()) for item in items]

        ignore = 'IGNORE' if ignore_duplicate else ''
        sql = f"INSERT {ignore} INTO {table} ({columns}) VALUES ({placeholders})"

        # 执行批量插入
        result = self.execute(
            sql,
            params=values if is_bulk else values[0],
            commit=commit
        )

        return result if is_bulk else result[0]

    def update(self, table: str, updates: Dict[str, Any], where: Dict[str, Any], commit: bool = True) -> int:
        """执行更新操作

        Args:
            table: 目标表名称
            updates: 更新字段字典（字段名: 新值）
            where: 条件字典（字段名: 期望值）
            commit: 是否自动提交事务

        Returns:
            本次更新影响的行数
        """
        set_clause = ', '.join([f"{k} = %s" for k in updates])
        where_clause = ' AND '.join([f"{k} = %s" for k in where])
        sql = f"UPDATE {table} SET {set_clause} WHERE {where_clause}"
        params = list(updates.values()) + list(where.values())
        return self.execute(sql, params=params, commit=commit)

    def delete(self, table: str, where: Dict[str, Any], commit: bool = True) -> int:
        """执行安全删除操作

        Args:
            table: 目标表名称
            where: 条件字典（字段名: 期望值）
            commit: 是否自动提交事务

        Returns:
            本次删除影响的行数
        """
        where_clause = ' AND '.join([f"{k} = %s" for k in where])
        sql = f"DELETE FROM {table} WHERE {where_clause}"

        return self.execute(sql, params=list(where.values()), commit=commit)

    def bulk_insert(self, table: str, columns: List[str], data: List[tuple], batch_size: int = 1000,
                    commit: bool = True) -> List[int]:
        """执行高性能批量插入

        Args:
            table: 目标表名称
            columns: 字段名称列表
            data: 元组数据列表（需与字段顺序一致）
            batch_size: 每批次插入量（默认1000）
            commit: 是否自动提交事务

        Returns:
            插入行ID列表（按插入顺序排列）

        Note:
            建议批量插入时关闭自动提交以提升性能
        """
        placeholders = ', '.join(['%s'] * len(columns))
        sql = f"INSERT INTO {table} ({', '.join(columns)}) VALUES ({placeholders})"

        ids = []
        for i in range(0, len(data), batch_size):
            batch = data[i:i + batch_size]
            ids.extend(self.execute(sql, params=batch, commit=False))

        if commit:
            self.conn.commit()

        return ids

    def __enter__(self):
        self.conn = pymysql.connect(host=self.host, port=self.port, user=self.user, password=self.password
                                    , db=self.db, charset='utf8')
        self.cursor = self.conn.cursor()
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        if hasattr(self, 'cursor'):
            print('关闭游标')
            self.cursor.close()
        if hasattr(self, 'conn'):
            print('关闭连接')
            self.conn.close()


class DbConnectionFactory:
    """数据库连接器工厂类，统一创建不同数据库类型的连接实例
    """
    _strategies: Dict[str, Type[BaseConnector]] = {
        'sqlite': SqlLiteConnector,
        'mysql': MysqlConnector
    }

    def create_connection(self, connection_type: str, **kwargs) -> BaseConnector:
        """创建指定类型的数据库连接器

        Args:
            connection_type: 数据库类型标识符，当前支持：
                - 'sqlite': SQLite数据库
                - 'mysql': MySQL数据库
            **kwargs: 连接参数，根据不同数据库类型传递不同参数：
                - SQLite: db_path (数据库文件路径)
                - MySQL: host, port, usr, pwd, db 等

        Returns:
            初始化完成的数据库连接器实例

        Raises:
            ValueError: 当指定不支持的数据库类型时抛出

        示例：
            >>> factory = DbConnectionFactory()
            >>> sqlite_conn = factory.create_connection('sqlite', db_path='test.db')
            >>> mysql_conn = factory.create_connection('mysql', host='localhost', db='mydb')
        """
        connector_cls = self._strategies.get(connection_type)
        if not connector_cls:
            raise ValueError(f"不支持的清理类型: {connection_type}")
        return connector_cls(**kwargs)

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        pass

#
# if __name__ == '__main__':
#     d = DbConnectionFactory()
#     d_obj = d.create_connection('sqlite',
#                                 db_path=r'/Users/makaiqiang/work/autoTest/qa/auto_test_result_db.sqlite').connect()
#     d_obj.insert('test_result_tb',
#                  {"id": 1, "execution_result": "pass", "title": "用例", "feature": "订单", "picture": "123",
#                   "picture_name": "222", "node_id": "123"})
#     print(d_obj.select('test_result_tb', '*'))
if __name__ == '__main__':
    import os
    from datetime import datetime
    
    # 定义数据库路径 - 使用项目正确的数据库文件
    db_path = r'F:\YQN202509\pw_test\auto_test_result_db.sqlite'
    
    # 确保目录存在
    db_dir = os.path.dirname(db_path)
    if not os.path.exists(db_dir):
        os.makedirs(db_dir)
        print(f"✅ 已创建目录: {db_dir}")
    
    # 创建数据库连接
    d = DbConnectionFactory()
    d_obj = d.create_connection('sqlite', db_path=db_path).connect()
    
    # 创建测试结果表(完整版本 - 使用"场景"字段替代node_id)
    create_table_sql = """
    CREATE TABLE IF NOT EXISTS test_result_tb (
        id INTEGER PRIMARY KEY AUTOINCREMENT,
        execution_result TEXT NOT NULL,
        title TEXT NOT NULL,
        feature TEXT,
        scenario TEXT NOT NULL,
        create_date TEXT NOT NULL,
        screenshot INTEGER DEFAULT 0,
        picture BLOB,
        picture_name TEXT
    )
    """
    d_obj.execute(create_table_sql)
    print(f"✅ 表 test_result_tb 已创建/已存在(使用场景字段)")
    print(f"💾 数据库位置: {db_path}")
    print(f"📁 文件存在: {os.path.exists(db_path)}")
    
    # 测试插入一条数据验证表结构
    d_obj.insert('test_result_tb',
                 {"execution_result": "pass", "title": "测试用例", "feature": "订单管理",
                  "scenario": "全字段搜索", "create_date": datetime.now().strftime('%Y-%m-%d')})
    print("✅ 测试数据插入成功")
    
    # 查询数据验证
    result = d_obj.select('test_result_tb',
                          ['id', 'execution_result', 'title', 'feature', 'scenario', 'create_date'])
    print(f"📊 查询结果: {result}")

